import torch
import torch.nn as nn
import torchvision.models as torch_models
import torchvision.transforms as transforms
from PIL import Image
import logging
import numpy as np

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class ImageRecognition:
    def __init__(self, model_path, num_classes):
        self.model_path = model_path
        self.num_classes = num_classes
        self.model = None
        self.feature_extractor = None
        self.transform = None
        self.device = None
        
    def initialize(self):
        """初始化模型和特征提取器"""
        try:
            logger.info('正在加载模型...')
            self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
            
            # 初始化模型
            self.model = torch_models.resnet50()
            self.model.fc = nn.Linear(self.model.fc.in_features, self.num_classes)
            
            # 添加安全设置
            torch.serialization.add_safe_globals([getattr])
            
            # 加载模型权重
            checkpoint = torch.load(self.model_path, map_location='cpu', weights_only=False)
            self.model.load_state_dict(checkpoint['state_dict'])
            self.model = self.model.to(self.device)
            self.model.eval()
            
            # 初始化特征提取器
            self.feature_extractor = nn.Sequential(*list(self.model.children())[:-1])
            self.feature_extractor.eval()
            
            # 设置图像预处理
            self.transform = transforms.Compose([
                transforms.Resize(256),
                transforms.CenterCrop(224),
                transforms.ToTensor(),
                transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                  std=[0.229, 0.224, 0.225])
            ])
            
            logger.info(f'模型加载成功，输出类别数: {self.num_classes}')
            return True
            
        except Exception as e:
            logger.error(f'初始化模型出错: {str(e)}')
            return False
            
    def extract_features(self, image_path):
        """从图片中提取特征向量"""
        try:
            # 加载并预处理图片
            image = Image.open(image_path)
            if image.mode != 'RGB':
                image = image.convert('RGB')
                
            # 使用预处理转换图片
            image_tensor = self.transform(image).unsqueeze(0).to(self.device)
            
            # 使用特征提取器获取特征
            with torch.no_grad():
                features = self.feature_extractor(image_tensor)
                features = features.cpu().numpy().flatten()
                
            return features
            
        except Exception as e:
            logger.error(f"特征提取失败: {str(e)}")
            return None
            
    def predict(self, image_path):
        """预测图片类别"""
        try:
            # 加载并预处理图片
            image = Image.open(image_path)
            image_tensor = self.transform(image).unsqueeze(0).to(self.device)
            
            # 提取特征
            with torch.no_grad():
                features = self.feature_extractor(image_tensor)
                features = features.squeeze().cpu().numpy()
            
            # 进行预测
            with torch.no_grad():
                outputs = self.model(image_tensor)
                probabilities = torch.nn.functional.softmax(outputs, dim=1)
                predicted_idx = torch.argmax(probabilities, dim=1).item()
                confidence = probabilities[0][predicted_idx].item()
            
            return {
                'success': True,
                'predicted_idx': predicted_idx,
                'confidence': float(confidence * 100),
                'features': features
            }
            
        except Exception as e:
            logger.error(f'预测出错: {str(e)}')
            return {
                'success': False,
                'error': str(e)
            }

    def predict_from_tensor(self, image_tensor):
        """从图像张量进行预测"""
        try:
            with torch.no_grad():
                outputs = self.model(image_tensor)
                _, predicted = outputs.max(1)
                predicted_idx = predicted.item()
                confidence = torch.softmax(outputs, dim=1)[0][predicted_idx].item() * 100
                
            return {
                'success': True,
                'predicted_idx': predicted_idx,
                'confidence': float(confidence)
            }
        except Exception as e:
            logger.error(f"从张量预测出错: {str(e)}")
            return {
                'success': False,
                'error': str(e)
            }